package first.group.ssh.channel;

import first.group.ssh.ConnectionCallback;
import first.group.ssh.PacketConnectStage;
import first.group.ssh.packet.ChannelClosePacket;
import first.group.ssh.packet.ChannelDataPacket;
import first.group.ssh.packet.ChannelEofPacket;
import first.group.ssh.packet.ChannelExtendedDataPacket;
import first.group.ssh.packet.ChannelFailurePacket;
import first.group.ssh.packet.ChannelOpenComfirmationPacket;
import first.group.ssh.packet.ChannelOpenFailurePacket;
import first.group.ssh.packet.ChannelRequestPacket;
import first.group.ssh.packet.ChannelSuccessPacket;
import first.group.ssh.packet.ChannelWindowAdjustPacket;
import first.group.ssh.packet.Packet;
import java.io.IOException;
import java.util.HashMap;
import java.util.logging.Level;
import java.util.logging.Logger;

public class ChannelMultiplexer extends PacketConnectStage {

    private static final Logger logger = Logger.getLogger("ChannelMultiplexer");

    private final ConnectionCallback connectionCallback;
    private final HashMap<Integer, Channel> channels = new HashMap<>();
    private int newChannelId = 0;

    public ChannelMultiplexer(ConnectionCallback connectionCallback) {
        this.connectionCallback = connectionCallback;
    }

    public void openChannel(Channel channel) {
        channel.open(this);
    }

    protected void addChannel(Channel channel) {
        var newId = newChannelId++;
        channels.put(newId, channel);
        channel.init(connectionCallback, newId);
    }

    protected void finishChannel(Channel channel) {
        channels.remove(channel.getLocalId());
    }

    private Channel getChannel(int localId) throws IOException {
        var channel = channels.get(localId);
        if (channel == null) {
            throw new IOException("Unknown local channel " + localId);
        }
        return channel;
    }

    @Override
    public void read(Packet packet) throws IOException {
        if (packet instanceof ChannelOpenComfirmationPacket comfirmationPacket) {
            var channel = getChannel(comfirmationPacket.recipientChannel);
            channel.openSucceed(
                    comfirmationPacket.senderChannel,
                    comfirmationPacket.initialWindowSize,
                    comfirmationPacket.maximumPacketSize,
                    comfirmationPacket.payload
            );
        } else if (packet instanceof ChannelOpenFailurePacket failurePacket) {
            var channel = getChannel(failurePacket.recipientChannel);
            channel.openFailed();
        } else if (packet instanceof ChannelDataPacket dataPacket) {
            var channel = getChannel(dataPacket.recipientChannel);
            channel.receiveData(dataPacket.data);
        } else if (packet instanceof ChannelExtendedDataPacket extendedDataPacket) {
            var channel = getChannel(extendedDataPacket.recipientChannel);
            channel.receiveData(extendedDataPacket.data);
        } else if (packet instanceof ChannelClosePacket closePacket) {
            var channel = getChannel(closePacket.recipientChannel);
            channel.receiveClose();
        } else if (packet instanceof ChannelEofPacket eofPacket) {
            var channel = getChannel(eofPacket.recipientChannel);
            channel.remoteEof();
        } else if (packet instanceof ChannelSuccessPacket successPacket) {
            var channel = getChannel(successPacket.recipientChannel);
            channel.requestSucceed();
        } else if (packet instanceof ChannelFailurePacket failurePacket) {
            var channel = getChannel(failurePacket.recipientChannel);
            channel.requestFailed();
        } else if (packet instanceof ChannelRequestPacket requestPacket) {
            var channel = getChannel(requestPacket.recipientChannel);
            if (requestPacket.request == null) {
                logger.log(Level.WARNING, "Unknown channel request: {0}", requestPacket.requestType);
                return;
            }
            channel.receiveRequestPacket(requestPacket.wantReply, requestPacket.request);
        } else if (packet instanceof ChannelWindowAdjustPacket windowAdjustPacket) {
            var channel = getChannel(windowAdjustPacket.recipientChannel);
            channel.enlargeRemoteWindowSize(windowAdjustPacket.bytesToAdd);
        } else {
            throw new AssertionError("Unknown packet: " + packet);
        }
    }

    public boolean shouldWrite() {
        for (var entry : channels.entrySet()) {
            if (entry.getValue().shouldWrite()) {
                return true;
            }
        }
        return false;
    }

    @Override
    public boolean shouldRead() {
        return true;
    }
}
